Support Mooncake 0.6 forward mode in find_alpha rules - #480
Merged
Conversation
Mooncake 0.6 reworks batched forward AD around `Lifted`/`NDual` and routes `prepare_derivative_cache` (used by `AutoMooncakeForward`) through `build_frule`. Add `frule!!(::Lifted, ...)` rules for `find_alpha` (float and integer third argument), gated behind `pkgversion(Mooncake) >= v"0.6"`, keeping the existing `Mooncake.Dual` rules for 0.5. Under 0.6 these `Lifted` rules intercept `find_alpha` as a primitive, so the batched pass never enters its body. Under 0.5 the batched (Nfwd) pass instead propagates `NDual`s straight through `find_alpha`, tripping on `ceil(::Int, ::NDual)` inside the root-finder; the `find_alpha(::NDual, ...)` short-circuit methods fix that and are therefore gated to `< v"0.6"` (verified redundant on 0.6, required on 0.5). This un-breaks the `AutoMooncakeForward` PlanarLayer-inverse integration case, so drop its `@test_broken` marker. Bump Mooncake compat to allow 0.6. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
yebai
force-pushed
the
compat-mooncake-0.6
branch
from
June 18, 2026 13:02
36852af to
c12c22f
Compare
Contributor
|
Bijectors.jl documentation for PR #480 is available at: |
Member
Should we wait here until Mooncake 0.6 release? |
Member
Author
Probably not, as this branch is fully compitable with 0.5, and the upcoming 0.6 is unlikely to change |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Mooncake 0.6 reworks batched forward AD around
Lifted/NDualand routesprepare_derivative_cache(used byAutoMooncakeForward) throughbuild_frule. Addfrule!!(::Lifted, ...)rules forfind_alpha(float and integer third argument), gated behindpkgversion(Mooncake) >= v"0.6", keeping the existingMooncake.Dualrules for 0.5.Under 0.6, these
Liftedrules interceptfind_alphaas a primitive, so the batched pass never enters its body. Under 0.5 the batched (Nfwd) pass instead propagatesNDuals straight throughfind_alpha, tripping onceil(::Int, ::NDual)inside the root-finder; thefind_alpha(::NDual, ...)short-circuit methods fix that and are therefore gated to< v"0.6"(verified redundant on 0.6, required on 0.5).This unbreaks the
AutoMooncakeForwardPlanarLayer-inverse integration case, so drop its@test_brokenmarker. Bump Mooncake compat to allow 0.6.Note: Mooncake 0.6 is not yet released. This PR fix Mooncake 0.6 compatibility in a non-breaking way.